package ru.ifmo.hmm;

import java.util.Arrays;
import java.util.Scanner;

/**
 * @author Iskander Akishev
 */
public class Util {

	public static class Model {
		public final double[][] a;
		public final double[][] b;
		public final double[] p0;
		public final int m;
		public final int n;

		public Model(double[][] a, double[][] b, double[] p0) {
			this.n = a.length;
			for (double[] ai : a) {
				if (ai.length != n) {
					throw new IllegalArgumentException();
				}
			}
			if ((p0.length != n) || (b.length != n)) {
				throw new IllegalArgumentException();
			}
			this.m = b[0].length;
			for (double[] bi : b) {
				if (bi.length != m) {
					throw new IllegalArgumentException();
				}
			}
			this.a = a;
			this.b = b;
			this.p0 = p0;
		}
	}

	public static Model readModel(Scanner in) {
		int n = in.nextInt();
		int m = in.nextInt();
		double[][] a = new double[n][n];
		double[][] b = new double[n][m];
		double[] p0 = new double[n];
		for (int i = 0; i < n; i++) {
			for (int j = 0; j < n; j++) {
				a[i][j] = in.nextDouble();
			}
		}
		for (int i = 0; i < n; i++) {
			for (int j = 0; j < m; j++) {
				b[i][j] = in.nextDouble();
			}
		}
		for (int i = 0; i < n; i++) {
			p0[i] = in.nextDouble();
		}
		return new Model(a, b, p0);
	}
	
	public static double dataProbability(Model m, int[] d) {
		int n = m.n;
		int k = d.length;
		double[][] p = new double[k][n];
		calcAlpha(p, m, d);

		double ans = 0d;
		for (int i = 0; i < n; i++) {
			ans += p[k - 1][i];
		}
		return ans;
	}

	public static int mostProbableState(Model m, int[] d, int t) {
		int k = d.length;
		double[][] alpha = new double[k][m.n];
		double[][] beta = new double[k][m.n];
		calcAlpha(alpha, m, d);
		calcBeta(beta, m, d);

		double maxVal = Double.NEGATIVE_INFINITY;
		int maxArg = -1;
		for (int i = 0; i < m.n; i++) {
			double cur = alpha[t][i] * beta[t][i];
			if (cur > maxVal) {
				maxVal = cur;
				maxArg = i;
			}
		}

		return maxArg;
	}

	public static int[] mostProbableSequence(Model m, int[] d) {
		int n = m.n;
		int k = d.length;
		double[][] gamma = new double[k][n];
		int[][] from = new int[k][n];

		for (int i = 0; i < n; i++) {
			gamma[0][i] = m.p0[i] * m.b[i][d[0]];
			from[0][i] = -1;
		}

		for (int j = 1; j < k; j++) {
			for (int i1 = 0; i1 < n; i1++) {
				double maxVal = Double.NEGATIVE_INFINITY;
				int maxArg = -1;
				for (int i2 = 0; i2 < n; i2++) {
					double cur = gamma[j - 1][i2] * m.a[i2][i1];
					if (cur > maxVal) {
						maxVal = cur;
						maxArg = i2;
					}
				}
				gamma[j][i1] = maxVal * m.b[i1][d[j]];
				from[j][i1] = maxArg;
			}

		}

		int i = 0;
		for (int z = 0; z < n; z++) {
			if (gamma[k - 1][z] > gamma[k - 1][i]) {
				i = z;
			}
		}
		int[] ans = new int[k];
		for (int j = k - 1; j >= 0; j--) {
			ans[j] = i;
			i = from[j][i];
		}
		return ans;
	}

	public static Model mostProbableModel(int[] d, int n, double eps) {
		int m = 0;
		for (int v : d) {
			if (v >= m) {
				m = v + 1;
			}
		}
		int k = d.length;

		double[][] a0 = new double[n][n];
		double[][] b0 = new double[n][m];
		double[] p0 = new double[n];
		for (int i = 0; i < n; i++) {
			Arrays.fill(a0[i], 1d / n);
			Arrays.fill(b0[i], 1d / m);
		}
		Arrays.fill(p0, 1d / n);
		Model model = new Model(a0, b0, p0);

		double[][][] xi = new double[k - 1][n][n];
		double[][] gamma = new double[k][n];
		double[][] alpha = new double[k][n];
		double[][] beta = new double[k][n];

		double curProb = dataProbability(model, d);
		while (true) {
			// Calc alpha, beta, xi & gamma
			calcAlpha(alpha, model, d);
			calcBeta(beta, model, d);
			for (int j = 0; j < k - 1; j++) {
				for (int i1 = 0; i1 < k; i1++) {
					for (int i2 = 0; i2 < k; i2++) {
						xi[j][i1][i2] = alpha[j][i1] * model.a[i1][i2] * model.b[i2][d[j + 1]] * beta[j + 1][i2] / curProb;
					}
				}
			}
			for (int j = 0; j < k; j++) {
				for (int i = 0; i < n; i++) {
					gamma[j][i] = alpha[j][i] * beta[j][i] / curProb;
				}
			}

			// Update model
			double[] z = new double[n];
			for (int i = 0; i < n; i++) {
				model.p0[i] = gamma[0][i];
				z[i] = 0d;
				for (int j = 0; j < k; j++) {
					z[i] += gamma[j][i];
				}
			}
			for (int i1 = 0; i1 < n; i1++) {
				for (int i2 = 0; i2 < n; i2++) {
					double v = 0d;
					for (int j = 0; j < k - 1; j++) {
						v += xi[j][i1][i2];
					}
					model.a[i1][i2] = v / z[i1];
				}
			}
			for (int i = 0; i < n; i++) {
				Arrays.fill(model.b[i], 0d);
			}

			for (int i = 0; i < n; i++) {
				for (int j = 0; j < k; j++) {
					model.b[i][d[j]] += gamma[j][i];
				}
			}
			for (int i = 0; i < n; i++) {
				for (int v = 0; v < m; v++) {
					model.b[i][v] /= z[i];
				}
			}

			// Calculate probability
			double newProb = dataProbability(model, d);
			assert (newProb >= curProb);
			if ((newProb - curProb) < eps) {
				break;
			}
			curProb = newProb;
		}
		return model;
	}

	private static void calcAlpha(double[][] p, Model m, int[] d) {
		int n = m.n;
		int k = d.length;

		if (p.length != k) {
			throw new IllegalArgumentException();
		}
		for (double[] pk : p) {
			if (pk.length != n) {
				throw new IllegalArgumentException();
			}
		}
		for (int i = 0; i < n; i++) {
			p[0][i] = m.p0[i] * m.b[i][d[0]];
		}

		for (int j = 1; j < k; j++) {
			for (int i1 = 0; i1 < n; i1++) {
				double v = 0d;
				for (int i2 = 0; i2 < n; i2++) {
					v += p[j - 1][i2] * m.a[i2][i1] * m.b[i1][d[j]];
				}
				p[j][i1] = v;
			}
		}
	}

	private static double[][] calcBeta(double[][] p, Model m, int[] d) {
		int n = m.n;
		int k = d.length;

		if (p.length != k) {
			throw new IllegalArgumentException();
		}
		for (double[] pk : p) {
			if (pk.length != n) {
				throw new IllegalArgumentException();
			}
		}
		for (int i = 0; i < n; i++) {
			p[k - 1][i] = 1d;
		}

		for (int j = k - 2; j >= 0; j--) {
			for (int i1 = 0; i1 < n; i1++) {
				double v = 0d;
				for (int i2 = 0; i2 < n; i2++) {
					v += p[j + 1][i2] * m.a[i1][i2] * m.b[i2][d[j]];
				}
				p[j][i1] = v;
			}
		}
		return p;
	}

}

